import asyncio

import dspy

# lm = dspy.LM("gemini/gemini-2.5-flash-preview-04-17", cache=False)
lm = dspy.LM("databricks/databricks-claude-3-7-sonnet", cache=False)
# lm = dspy.LM("openai/gpt-4o-mini", cache=False)
dspy.settings.configure(lm=lm, adapter=dspy.XMLAdapter())


class MyModule(dspy.Module):
    def __init__(self):
        super().__init__()

        self.predict1 = dspy.Predict("question->answer")
        self.predict2 = dspy.Predict("question, answer->judgement, score")

    def forward(self, question: str, **kwargs):
        answer = self.predict1(question=question)
        judgement = self.predict2(question=question, answer=answer)
        return judgement


predict = MyModule()
stream_listeners = [
    dspy.streaming.StreamListener(signature_field_name="answer"),
    dspy.streaming.StreamListener(signature_field_name="judgement"),
]
stream_predict = dspy.streamify(
    predict,
    stream_listeners=stream_listeners,
    include_final_prediction_in_output_stream=False,
)


async def main():
    with dspy.settings.context(adapter=dspy.XMLAdapter()):
        output = stream_predict(question="why did a chicken cross the kitchen?")

        return_value = None
        async for value in output:
            if isinstance(value, dspy.streaming.StatusMessage):
                print(value.message)
            elif isinstance(value, dspy.Prediction):
                print(value)
            elif isinstance(value, dspy.streaming.StreamResponse):
                print(value)
            else:
                print(value.choices[0].delta.content)
    return return_value


import time

start = time.time()
asyncio.run(main())
end = time.time()


print(f"Time taken: {end - start} seconds")

dspy.inspect_history()